// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/connectors/connector.h>

namespace kumo::pollux::connector {
    namespace {
        std::unordered_map<std::string, std::shared_ptr<ConnectorFactory> > &
        connectorFactories() {
            static std::unordered_map<std::string, std::shared_ptr<ConnectorFactory> >
                    factories;
            return factories;
        }

        std::unordered_map<std::string, std::shared_ptr<Connector> > &connectors() {
            static std::unordered_map<std::string, std::shared_ptr<Connector> > connectors;
            return connectors;
        }
    } // namespace

    bool DataSink::Stats::empty() const {
        return numWrittenBytes == 0 && numWrittenFiles == 0 && spillStats.empty();
    }

    std::string DataSink::Stats::toString() const {
        return fmt::format(
            "numWrittenBytes {} numWrittenFiles {} {}",
            succinctBytes(numWrittenBytes),
            numWrittenFiles,
            spillStats.toString());
    }

    bool registerConnectorFactory(std::shared_ptr<ConnectorFactory> factory) {
        bool ok =
                connectorFactories().insert({factory->connectorName(), factory}).second;
        POLLUX_CHECK(
            ok,
            "ConnectorFactory with name '{}' is already registered",
            factory->connectorName());
        return true;
    }

    bool hasConnectorFactory(const std::string &connectorName) {
        return connectorFactories().count(connectorName) == 1;
    }

    bool unregisterConnectorFactory(const std::string &connectorName) {
        auto count = connectorFactories().erase(connectorName);
        return count == 1;
    }

    std::shared_ptr<ConnectorFactory> getConnectorFactory(
        const std::string &connectorName) {
        auto it = connectorFactories().find(connectorName);
        POLLUX_CHECK(
            it != connectorFactories().end(),
            "ConnectorFactory with name '{}' not registered",
            connectorName);
        return it->second;
    }

    bool registerConnector(std::shared_ptr<Connector> connector) {
        bool ok = connectors().insert({connector->connectorId(), connector}).second;
        POLLUX_CHECK(
            ok,
            "Connector with ID '{}' is already registered",
            connector->connectorId());
        return true;
    }

    bool unregisterConnector(const std::string &connectorId) {
        auto count = connectors().erase(connectorId);
        return count == 1;
    }

    std::shared_ptr<Connector> getConnector(const std::string &connectorId) {
        auto it = connectors().find(connectorId);
        POLLUX_CHECK(
            it != connectors().end(),
            "Connector with ID '{}' not registered",
            connectorId);
        return it->second;
    }

    const std::unordered_map<std::string, std::shared_ptr<Connector> > &
    getAllConnectors() {
        return connectors();
    }

    melon::Synchronized<
        std::unordered_map<std::string_view, std::weak_ptr<cache::ScanTracker> > >
    Connector::trackers_;

    // static
    void Connector::unregisterTracker(cache::ScanTracker *tracker) {
        trackers_.withWLock([&](auto &trackers) { trackers.erase(tracker->id()); });
    }

    std::shared_ptr<cache::ScanTracker> Connector::getTracker(
        const std::string &scanId,
        int32_t loadQuantum) {
        return trackers_.withWLock([&](auto &trackers) -> auto {
            auto it = trackers.find(scanId);
            if (it == trackers.end()) {
                auto newTracker = std::make_shared<cache::ScanTracker>(
                    scanId, unregisterTracker, loadQuantum);
                trackers[newTracker->id()] = newTracker;
                return newTracker;
            }
            std::shared_ptr<cache::ScanTracker> tracker = it->second.lock();
            if (!tracker) {
                tracker = std::make_shared<cache::ScanTracker>(
                    scanId, unregisterTracker, loadQuantum);
                trackers[tracker->id()] = tracker;
            }
            return tracker;
        });
    }

    std::string commitStrategyToString(CommitStrategy commitStrategy) {
        switch (commitStrategy) {
            case CommitStrategy::kNoCommit:
                return "NO_COMMIT";
            case CommitStrategy::kTaskCommit:
                return "TASK_COMMIT";
            default:
                POLLUX_UNREACHABLE(
                    "UNKOWN COMMIT STRATEGY: {}", static_cast<int>(commitStrategy));
        }
    }

    CommitStrategy stringToCommitStrategy(const std::string &strategy) {
        if (strategy == "NO_COMMIT") {
            return CommitStrategy::kNoCommit;
        } else if (strategy == "TASK_COMMIT") {
            return CommitStrategy::kTaskCommit;
        } else {
            POLLUX_UNREACHABLE("UNKOWN COMMIT STRATEGY: {}", strategy);
        }
    }

    melon::Dynamic ColumnHandle::serializeBase(std::string_view name) {
        melon::Dynamic obj = melon::Dynamic::object;
        obj["name"] = name;
        return obj;
    }

    melon::Dynamic ColumnHandle::serialize() const {
        return serializeBase("ColumnHandle");
    }

    melon::Dynamic ConnectorTableHandle::serializeBase(
        std::string_view name) const {
        melon::Dynamic obj = melon::Dynamic::object;
        obj["name"] = name;
        obj["connectorId"] = connectorId_;
        return obj;
    }

    melon::Dynamic ConnectorTableHandle::serialize() const {
        return serializeBase("ConnectorTableHandle");
    }
} // namespace kumo::pollux::connector
